#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <unistd.h>
#include <libgen.h>
#include <sys/stat.h>
#include <sys/select.h>

#define PORT 3456//服务器监听端口
#define SERVER_IP "127.0.0.1"
#define BLOCK_SIZE 512//TFTP数据块大小
#define BASE_DIR "/home/linux/tftp_files"//文件存储根目录

typedef enum{
    //读请求
    OP_RRQ = 1,
    //写请求
    OP_WRQ = 2,
    //数据包
    OP_DATA = 3,
    //确认包
    OP_ACK = 4,
    //错误包
    OP_ERROR = 5
}tftp_opcode;
//目录创建函数
void ensure_dir_exists(const char *path){
    struct stat st = {0};
    if(stat(path, &st) == -1){
        if(mkdir(path, 0755) == -1){
            perror("创建目录失败");
            exit(EXIT_FAILURE);
        }else{
            printf("目录 %s 创建成功\n", path);
        }
    }
}

void send_ack(int socket, struct sockaddr_in*addr, uint16_t block){
    uint8_t ack[4] = {0, OP_ACK, (block >> 8) & 0xFF, block & 0xFF};//构建ACK包
    //包头结构：2字节操作码（0x0004）+ 2字节块号（大端序）
    sendto(socket, ack, sizeof(ack), 0, (struct sockaddr*)addr, sizeof(*addr));//发送ACK包
}

//下载函数
void handle_download(int socket, struct sockaddr_in *client_addr, char *filename){
    printf("尝试下载文件: %s\n", filename);
    //提纯文件名（过滤路径，防止../遍历攻击）
    char *fname_copy = strdup(filename);//复制文件名(basename会修改传入的字符串)
    char *pure_filename = basename(fname_copy);//提纯文件名(“/a/b.txt -> "b.txt”)

    //构建存储路径 （BASE_DIR+纯文件名）
    char filepath[256];
    snprintf(filepath, sizeof(filepath), "%s/%s", BASE_DIR, pure_filename);//防止缓冲区溢出
    free(fname_copy);//释放复制的文件名内存
    printf("服务器文件路径： %s\n", filepath);

    int fd = open(filepath, O_RDONLY);//以只读方式打开文件
    if(fd < 0){
        perror("打开文件失败:");
        //构建错误包：操作5+错误码1+错误消息
        uint8_t err[5 + sizeof("文件未找到")] = {0, OP_ERROR, 1};
        strcpy((char*)&err[4], "文件未找到");//错误消息是从第五个字节开始
        sendto(socket, err, sizeof(err), 0, (struct sockaddr*)client_addr, sizeof(*client_addr));//发送错误包
        return;
    }

    uint8_t buf[BLOCK_SIZE + 4];//数据包缓冲区（512字节数据+4字节头部）
    uint16_t block = 1;//数据块编号从1开始
    ssize_t bytes_read;//实际读取字节数
    while((bytes_read = read(fd, &buf[4], BLOCK_SIZE)) > 0){
        //读取文件数据
        //构建数据包头部
        buf[0] = 0;//操作码高位（OP_DATA为3， 0x0003）
        buf[1] = OP_DATA;//操作码低位
        buf[2] = (block >> 8) & 0xFF;//块编号高位（大端序）
        buf[3] = block & 0xFF;//块编号低位

        printf("发送数据块 %d, 大小 %zd 字节\n", block, bytes_read);
        //发送数据包
        sendto(socket, buf, bytes_read + 4, 0, (struct sockaddr*)client_addr, sizeof(*client_addr));

        //设置5秒超时等待ACK
        fd_set read_fds;
        FD_ZERO(&read_fds);//清空文件描述符集合
        FD_SET(socket, &read_fds);//将socket加入集合
        struct timeval timeout = {5, 0};//设置5秒超时
        
        int ret = select(socket + 1, &read_fds, NULL, NULL, &timeout);//等待可读事件
        if(ret <= 0)//超时或错误
        {
            printf("等待ACK超时，重发数据块 %d\n", block);
            continue;//重发当前数据块
        }

        //接收ACK包 4字节
        ssize_t ack_size = recvfrom(socket, buf, 4, 0, NULL, NULL);
        //验证ack包有效性：操作码是否为ACK且块编号匹配
        if(ack_size < 4 || buf[1] != OP_ACK || ntohs(*(uint16_t*)&buf[2]) != block){
            printf("收到无效ACK，重发数据块 %d\n", block);
            continue;//重发当前数据块
        }

        printf("收到ACK %d\n", block);
        block++;//块编号递增
    }
    //发送结束包（空数据块，4字节包头，块号为最后一块 + 1）
    if(bytes_read == 0){
        buf[0] = 0;
        buf[1] = OP_DATA;
        buf[2] = (block >> 8) & 0xFF;
        buf[3] = block & 0xFF;
        sendto(socket, buf, 4, 0, (struct sockaddr*)client_addr, sizeof(*client_addr));//空包
        printf("文件传输完成，发送结束包\n");

        //等待客户端回应对结束包的ACK
        recvfrom(socket, buf, 4, 0, NULL, NULL);
    }
    close(fd);//关闭文件描述符
    printf("下载处理完成\n");
}
//上传处理函数
void handle_upload(int socket, struct sockaddr_in *client_addr, char *filename){
     printf("尝试上传文件: %s\n", filename);
    //提纯文件名（过滤路径，防止../遍历攻击）
    char *fname_copy = strdup(filename);//复制文件名(basename会修改传入的字符串)
    char *pure_filename = basename(fname_copy);//提纯文件名(“/a/b.txt -> "b.txt”)

    //构建存储路径 （BASE_DIR+纯文件名）
    char filepath[256];
    snprintf(filepath, sizeof(filepath), "%s/%s", BASE_DIR, pure_filename);//防止缓冲区溢出
    free(fname_copy);//释放复制的文件名内存

    //创建父目录（a/b/c.txt -> 创建a/b）
    char *dir = strdup(filepath);
    ensure_dir_exists(dirname(dir));
    free(dir);

    //创建文件
    int fd = open(filepath, O_WRONLY | O_TRUNC, 0644);//以写入方式创建文件
    if(fd < 0){
        perror("创建文件失败:");
        return;
    }
    uint8_t buf[BLOCK_SIZE + 4];//数据包缓冲区（512字节数据+4字节头部）
    uint16_t expected_block = 1;//期望的数据块编号从1开始

    send_ack(socket, client_addr, 0);//响应WRQ,发送ACK块号0（块号0，表示准备好接收）
    while(1){
        socklen_t len = sizeof(*client_addr);
        //接收客户端发送的数据包
        ssize_t nbytes = recvfrom(socket, buf, sizeof(buf), 0, (struct sockaddr*)client_addr, &len);
        if(nbytes < 4)//数据包过小
        {
            printf("收到无效数据包，忽略\n");
            break;
        }
        
        //处理数据块
        if(buf[1] == OP_DATA && ntohs(*(uint16_t*)&buf[2]) == expected_block){
            size_t data_size = nbytes - 4;//实际数据大小
            printf("收到块号 %d, 大小 %zu 字节\n", expected_block, data_size);
            //写入文件
            write(fd, &buf[4], data_size);
            send_ack(socket, client_addr, expected_block);//发送ACK
            //如果数据块不足512字节，视为最后一块
            if(data_size < BLOCK_SIZE){
                printf("文件传输完成\n");
                break;
            }
            expected_block++;//期望块号递增
        }else if(buf[1] == OP_ERROR){
            //处理错误包
            printf("错误： %s\n", &buf[4]);
            break;
        }else{
            printf("收到意外包：操作码：%d，块号%d\n", buf[1], ntohs(*(uint16_t*)&buf[2]));
    }
}
    close(fd);//关闭文件描述符
    printf("上传处理完成\n");
}

int main(){
    //确保基础存储目录存在
    printf("基础存储目录: %s\n", BASE_DIR);
    //调用目录创建函数
    ensure_dir_exists(BASE_DIR);
    //创建UDPsocket
    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    server_addr.sin_addr.s_addr = inet_addr(SERVER_IP);
    //绑定socket到指定端口与地址
    int ret_bind = bind(sockfd, (struct sockaddr*)&server_addr, sizeof(server_addr));
    if(ret_bind == -1){
        perror("绑定失败");
        close(sockfd);
        return -1;
    }
    printf("TFTP服务器已启动，监听端口 %d，服务器ip：%s\n", PORT, SERVER_IP);
    printf("等待客户端连接...\n");
    //开启主循环，等待客户端请求
    while(1){
        //接收缓冲区(最大的数据包512字节+4字节头部)
        uint8_t buffer[516];
        //客户端地址结构体
        struct sockaddr_in client_addr;
        socklen_t len = sizeof(client_addr);//地址长度初始化
        //接收客户端请求
        ssize_t nbytes = recvfrom(sockfd, buffer, sizeof(buffer), 0,(struct sockaddr*)&client_addr, &len);

        if(nbytes > 2){
            //返回至少包含了操作码（2字节）+文件名（1字节）
            char *filename = (char*)&buffer[2];//文件名从第3个字节开始
            char *mode = filename + strlen(filename) + 1;//固定模式字段（“octet”）
            char client_ip[INET_ADDRSTRLEN];//存储客户端IP地址字符串
            //将客户端ip转换为字符串
            inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
            printf("收到来自客户端%s:%d的请求\n",client_ip, ntohs(client_addr.sin_port));
            //判断操作码，并且去调用对应的处理函数
            if(buffer[1] == OP_RRQ){//下载请求
                printf("下载请求：%s\n", filename);
                //调用处理下载请求的函数
                handle_download(sockfd, &client_addr, filename); 
                printf("下载完成\n");
            }else if(buffer[1] == OP_WRQ){//上传请求
                printf("上传请求：%s\n", filename);
                //调用处理上传请求的函数
                handle_upload(sockfd, &client_addr, filename);
                printf("上传完成\n");
            }else{
                printf("未知请求类型：%d\n", buffer[1]);
            }
        }
    }

    close(sockfd);
    return 0;
}
